import os
import sys
from functools import wraps
from loguru import logger
import service.pb.model_server_pb2 as pb

CHUNK_SIZE = 1024 * 1024  # 1MB


def get_file_chunk(full_file_path):
    with open(full_file_path, "rb") as f:
        while True:
            chunks = f.read(CHUNK_SIZE)
            if len(chunks) == 0:
                return
            yield pb.FileChunk(buffer=chunks)


def save_chunks_to_file(chunks, full_file_path):
    with open(full_file_path, 'wb') as f:
        for chunk in chunks:
            f.write(chunk.buffer)


def debug_exit(func):
    @wraps(func)
    def debug_exit_decorator(*args, **kwargs):
        try:
            result = func(*args, **kwargs)
            return result
        except Exception as e:
            logger.exception("Exit func: {}".format(func.__name__))
            os.kill(os.getpid(), 15)
    return debug_exit_decorator
